from PIL import Image
from utils.call_llm import client
from utils.img_utils import image_to_base64
import os
from dotenv import load_dotenv

load_dotenv()

def extract_text_from_image(image: Image.Image, prompt: str = None) -> str:
    """Extract text from image using OpenAI Vision API
    
    Args:
        image (PIL.Image): Image to process
        prompt (str, optional): Custom prompt for extraction. Defaults to general OCR.
        
    Returns:
        str: Extracted text from image
    """
    # Convert image to base64
    img_base64 = image_to_base64(image)
    
    # Default prompt for general OCR
    if prompt is None:
        prompt = "Please extract all text from this image."
    
    try:
        model = os.environ.get("VERSION_MODEL", "")
        # Call Vision API
        response = client.chat.completions.create(
            model=model,
            messages=[{
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_base64}"}}
                ]
            }]
        )

        return response.choices[0].message.content
    except Exception as e:
        print(f"Error extracting text: {e}")
        import time
        time.sleep(60)
        # Call Vision API
        response = client.chat.completions.create(
            model=model,
            messages=[{
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_base64}"}}
                ]
            }]
        )

        return response.choices[0].message.content
if __name__ == "__main__":
    # Test vision processing
    test_image = Image.open("example.png")
    result = extract_text_from_image(test_image)
    print("Extracted text:", result)